import logging
import time

from train import disL2, reload


def load_data():
    entity, relation, G = {}, {}, []
    with open('FB15k/entity2id.txt', 'r') as f:
        for line in f:
            temp = line.replace(' ', '\t').strip().split('\t')
            assert len(temp) >= 2, 'decode ERROR'
            temp = list(map(str.strip, temp))
            entity[temp[0]] = temp[1]
    with open('FB15k/relation2id.txt', 'r') as f:
        for line in f:
            temp = line.replace(' ', '\t').strip().split('\t')
            assert len(temp) >= 2, 'decode ERROR'
            temp = list(map(str.strip, temp))
            relation[temp[0]] = temp[1]
    with open('FB15k/test.txt', 'r') as f:
        for line in f:
            temp = line.replace(' ', '\t').strip().split('\t')
            assert len(temp) >= 3, 'decode ERROR'
            temp = list(map(str.strip, temp))
            row = [entity[temp[0]], entity[temp[1]], relation[temp[2]]]
            G.append(row)
    return entity, relation, G


def do_test(entity, relation, G, entity_data, relation_data):
    rank, hits10 = 0, 0
    index = 1
    # print('ee',entity,G,entity_data,relation_data)
    for triple in G:
        head, tail = [], []
        for i in entity.values():
            h = int(i)
            t = int(triple[1])
            r = int(triple[2])
            head.append([h, t, r, disL2(entity_data[h], relation_data[r], entity_data[t])])
            h = int(triple[0])
            t = int(i)
            r = int(triple[2])
            tail.append([h, t, r, disL2(entity_data[h], relation_data[r], entity_data[t])])
        sorted_rank_head = sorted(head, key=lambda x: x[-1])
        sorted_rank_tail = sorted(tail, key=lambda x: x[-1])
        for i in range(len(sorted_rank_head)):
            if int(sorted_rank_head[i][0]) == int(triple[0]):
                if i < 10:
                    hits10 += 1
                rank += i + 1
                break
        for i in range(len(sorted_rank_tail)):
            if int(sorted_rank_tail[i][1]) == int(triple[1]):
                if i < 10:
                    hits10 += 1
                rank += i + 1
                break
        if index % 10 == 0:
            logging.info('{}/{} rank {:.2f}, hits {:.2f}%'.
                         format(index, len(G), rank / (2 * index), 100 * hits10 / (2 * index)))
        index += 1
    rank = rank / (2 * len(G))
    hits10 = hits10 / (2 * len(G))
    print('mean_rank', rank, 'hits10', hits10)
    return rank, hits10


if __name__ == '__main__':
    st = time.time()
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s - %(filename)s[line:%(lineno)d]: %(message)s')
    entity, relation, G = load_data()
    # 'FB15k/entity_out_ori.bin', 'FB15k/relation_out_ori.bin'
    # 'FB15k/entity_out_tf.bin', 'FB15k/relation_out_tf.bin'
    entity_data, relation_data = reload('FB15k/entity_out_tf.bin', 'FB15k/relation_out_tf.bin')
    # entity_data = list(entity_data.values())
    # relation_data = list(relation_data.values())
    do_test(entity, relation, G, entity_data, relation_data)
